
#ifndef __HOOKD3D_H__
#define __HOOKD3D_H__

#include <d3d9.h>
#include <D3D10_1.h>
//#include <d3dx10.h>

template <class t, int offset>
class a
{
	a(){*(LPDWORD)((m_pI->lpVtbl) + offset) = NULL;}

	t* m_pI;
};


//class IDXGISwapChainDWM : public IDXGISwapChain { /*...*/ };
interface IDXGISwapChainDWM;

typedef struct IDXGISwapChainDWMVtbl
{
	BEGIN_INTERFACE

		HRESULT ( STDMETHODCALLTYPE *QueryInterface )( 
		IDXGISwapChain * This,
		/* [in] */ REFIID riid,
		/* [annotation][iid_is][out] */ 
		__RPC__deref_out  void **ppvObject);

		ULONG ( STDMETHODCALLTYPE *AddRef )( 
			IDXGISwapChain * This);

		ULONG ( STDMETHODCALLTYPE *Release )( 
			IDXGISwapChainDWM * This);

		HRESULT ( STDMETHODCALLTYPE *SetPrivateData )( 
			IDXGISwapChain * This,
			/* [in] */ REFGUID Name,
			/* [in] */ UINT DataSize,
			/* [in] */ const void *pData);

		HRESULT ( STDMETHODCALLTYPE *SetPrivateDataInterface )( 
			IDXGISwapChain * This,
			/* [in] */ REFGUID Name,
			/* [in] */ const IUnknown *pUnknown);

		HRESULT ( STDMETHODCALLTYPE *GetPrivateData )( 
			IDXGISwapChain * This,
			/* [in] */ REFGUID Name,
			/* [out][in] */ UINT *pDataSize,
			/* [out] */ void *pData);

		HRESULT ( STDMETHODCALLTYPE *GetParent )( 
			IDXGISwapChain * This,
			/* [in] */ REFIID riid,
			/* [retval][out] */ void **ppParent);

		HRESULT ( STDMETHODCALLTYPE *GetDevice )( 
			IDXGISwapChain * This,
			/* [in] */ REFIID riid,
			/* [retval][out] */ void **ppDevice);

		HRESULT ( STDMETHODCALLTYPE *Present )( 
			IDXGISwapChain * This,
			/* [in] */ UINT SyncInterval,
			/* [in] */ UINT Flags);

		HRESULT ( STDMETHODCALLTYPE *GetBuffer )( 
			IDXGISwapChain * This,
			/* [in] */ UINT Buffer,
			/* [in] */ REFIID riid,
			/* [out][in] */ void **ppSurface);

		HRESULT ( STDMETHODCALLTYPE *GetDesc )( 
			IDXGISwapChain * This,
			/* [out] */ DXGI_SWAP_CHAIN_DESC *pDesc);

		HRESULT ( STDMETHODCALLTYPE *ResizeBuffers )( 
			IDXGISwapChain * This,
			/* [in] */ UINT BufferCount,
			/* [in] */ UINT Width,
			/* [in] */ UINT Height,
			/* [in] */ DXGI_FORMAT NewFormat,
			/* [in] */ UINT SwapChainFlags);

		HRESULT ( STDMETHODCALLTYPE *ResizeTarget )( 
			IDXGISwapChain * This,
			/* [in] */ const DXGI_MODE_DESC *pNewTargetParameters);

		HRESULT ( STDMETHODCALLTYPE *GetContainingOutput )( 
			IDXGISwapChain * This,
			IDXGIOutput **ppOutput);

		HRESULT ( STDMETHODCALLTYPE *GetFrameStatistics )( 
			IDXGISwapChain * This,
			/* [out] */ DXGI_FRAME_STATISTICS *pStats);

		HRESULT ( STDMETHODCALLTYPE *GetLastPresentCount )( 
			IDXGISwapChain * This,
			/* [out] */ UINT *pLastPresentCount);
		HRESULT ( STDMETHODCALLTYPE *SetFullscreenState )( 
			IDXGISwapChain * This,
			/* [in] */ BOOL Fullscreen,
			/* [in] */ IDXGIOutput *pTarget);

		HRESULT ( STDMETHODCALLTYPE *GetFullscreenState )( 
			IDXGISwapChain * This,
			/* [out] */ BOOL *pFullscreen,
			/* [out] */ IDXGIOutput **ppTarget);

	END_INTERFACE
} IDXGISwapChainDWMVtbl;

interface IDXGISwapChainDWM
{
	CONST_VTBL struct IDXGISwapChainDWMVtbl *lpVtbl;
};

template<class iface, int offset, typename typefunc>
class CComHooker
{
public:
	typedef typefunc *PPROC;
	typedef iface IFACE_TYPE;
private:
	BOOL m_bHooked;
	PPROC m_pfnNew;
	iface* m_pI;
	PPROC m_pfnOld;
public:
	operator PPROC()
	{
		return m_pfnOld;
	}

	CComHooker()
	{
	}

	CComHooker(PPROC pfnNew) : m_pfnNew(pfnNew), m_bHooked(FALSE), m_pI(NULL), m_pfnOld(NULL)
	{
	}

	~CComHooker()
	{
		Destroy();
	}

	operator iface*()
	{
		return m_pI;
	}

	void* GetVtbl()
	{
		return m_pI->lpVtbl;
	}

	PPROC& GetProcAddr()
	{
		return (PPROC&)(*((LPDWORD)(((LPBYTE)(m_pI->lpVtbl)) + offset)));
	}

	BOOL Create(iface* pI)
	{
		m_pI = pI;
		return TRUE;
	}

	BOOL IsHooked(){return m_bHooked;}
	void Destroy()
	{
		if (m_bHooked)
			Unhook();
		m_pI = NULL;
	}

	BOOL Hook();
	BOOL Unhook();
};

#define HOOK_FUNC(iface, method) iface##method
#define DEFINE_HOOK(iface, method) template class CComHooker<iface, FIELD_OFFSET(iface##Vtbl, method), HOOK_FUNC(iface, method)>;

#define CCOMHOOKER(iface, method) CComHooker_##iface##_##method

#define DECLARE_COMHOOKER(iface, method) \
class CCOMHOOKER(iface, method)\
{\
public:\
	typedef PROTOTYPE;\
	typedef PROC *PPROC;\
	typedef iface IFACE_TYPE;\
private:\
	BOOL m_bHooked;\
	PPROC m_pfnNew;\
	iface* m_pI;\
	PPROC m_pfnOld;\
public:\
	operator PPROC()\
	{\
		return m_pfnOld;\
	}\
	CCOMHOOKER(iface, method)()\
	{\
	}\
	CCOMHOOKER(iface, method)(PPROC pfnNew) : m_pfnNew(pfnNew), m_bHooked(FALSE), m_pI(NULL), m_pfnOld(NULL)\
	{\
	}\
	~CCOMHOOKER(iface, method)()\
	{\
		Destroy();\
	}\
	operator iface*()\
	{\
		return m_pI;\
	}\
	void* GetVtbl();\
	PPROC GetProcAddr();\
	BOOL Create(iface* pI);\
	BOOL IsHooked(){return m_bHooked;}\
	void Destroy()\
	{\
		if (m_bHooked)\
			Unhook();\
		m_pI = NULL;\
	}\
	BOOL Hook();\
	BOOL Unhook();\
};

#define DEFINE_COMHOOKER(iface, method) \
void* CCOMHOOKER(iface, method)::GetVtbl()\
{\
	return m_pI->lpVtbl;\
}\
\
CCOMHOOKER(iface, method)::PPROC CCOMHOOKER(iface, method)::GetProcAddr()\
{\
	return m_pI->lpVtbl->method;\
}\
\
BOOL CCOMHOOKER(iface, method)::Create(iface* pI)\
{\
	m_pI = pI;\
	return TRUE;\
}\
\
BOOL CCOMHOOKER(iface, method)::Hook()\
{\
	char szBuf[512];\
	SYSTEMTIME st;\
	DWORD dwProtectOld;\
	if (!VirtualProtect(&m_pI->lpVtbl->method, sizeof(m_pI->lpVtbl->method), PAGE_EXECUTE_READWRITE, &dwProtectOld))\
		return FALSE;\
	m_pfnOld = m_pI->lpVtbl->method;\
	m_pI->lpVtbl->method = m_pfnNew;\
	VirtualProtect(&m_pI->lpVtbl->method, sizeof(m_pI->lpVtbl->method), dwProtectOld, &dwProtectOld);\
	m_bHooked = TRUE;\
	GetLocalTime(&st);\
	sprintf(szBuf, "D3D Hook[%04d, %04d]"#iface"(%08p.%08p)"#method" hooked: %08p->%08p\n", GetCurrentProcessId(), GetCurrentThreadId(), m_pI, m_pI->lpVtbl, m_pfnOld, m_pfnNew);\
	OutputDebugStringA(szBuf);\
	return TRUE;\
}\
\
BOOL CCOMHOOKER(iface, method)::Unhook()\
{\
	char szBuf[512];\
	SYSTEMTIME st;\
	DWORD dwProtectOld;\
	if (!VirtualProtect(&m_pI->lpVtbl->method, sizeof(m_pI->lpVtbl->method), PAGE_EXECUTE_READWRITE, &dwProtectOld))\
		return FALSE;\
	m_pI->lpVtbl->method = m_pfnOld;\
	VirtualProtect(&m_pI->lpVtbl->method, sizeof(m_pI->lpVtbl->method), dwProtectOld, &dwProtectOld);\
	m_bHooked = FALSE;\
	GetLocalTime(&st);\
	sprintf(szBuf, "D3D Hook[%04d, %04d]"#iface"(%08p.%08p)"#method" unhooked\n", GetCurrentProcessId(), GetCurrentThreadId(), m_pI, m_pI->lpVtbl);\
	OutputDebugStringA(szBuf);\
	/*Sleep(10000);*/\
	return TRUE;\
}

#undef PROTOTYPE
#define PROTOTYPE HRESULT STDMETHODCALLTYPE PROC(IDirect3DSwapChain9* This, CONST RECT* pSourceRect,CONST RECT* pDestRect,HWND hDestWindowOverride,CONST RGNDATA* pDirtyRegion,DWORD dwFlags)

DECLARE_COMHOOKER(IDirect3DSwapChain9, Present)

typedef CCOMHOOKER(IDirect3DSwapChain9, Present) CD3dHookPresent;

#undef PROTOTYPE
#define PROTOTYPE HRESULT STDMETHODCALLTYPE PROC(IDirect3DDevice9* This, D3DPRIMITIVETYPE PrimitiveType,UINT StartVertex,UINT PrimitiveCount)

DECLARE_COMHOOKER(IDirect3DDevice9, DrawPrimitive)

typedef CCOMHOOKER(IDirect3DDevice9, DrawPrimitive) CD3D9HookDrawPrimitive;

#undef PROTOTYPE
#define PROTOTYPE HRESULT STDMETHODCALLTYPE PROC(IDXGISwapChain* This, UINT SyncInterval, UINT Flags);

DECLARE_COMHOOKER(IDXGISwapChainDWM, Present)
typedef CCOMHOOKER(IDXGISwapChainDWM, Present) CDXGIHookPresent;

#undef PROTOTYPE
#define PROTOTYPE void STDMETHODCALLTYPE PROC(ID3D10Device* This, UINT VertexCount, UINT StartVertexLocation);

DECLARE_COMHOOKER(ID3D10Device, Draw)
typedef CCOMHOOKER(ID3D10Device, Draw) CD3D10HookDraw;

#undef PROTOTYPE
#define PROTOTYPE void STDMETHODCALLTYPE PROC(ID3D10Device* This, UINT IndexCount, UINT StartIndexLocation, INT BaseVertexLocation);

DECLARE_COMHOOKER(ID3D10Device, DrawIndexed)
typedef CCOMHOOKER(ID3D10Device, DrawIndexed) CD3D10HookDrawIndexed;

#endif
